import torch
from model import PromptBert
import random
import os
import onnx
import onnxruntime
import numpy as np
import time


def pt2onnx():
    class Config:
        root_path = os.path.dirname(os.path.realpath(__file__))
        model_path = '/data0/jianyu10/PTM/huggingface_model_cache/chinese-roberta-wwm-ext'

        dropout_prob = 0.25
        device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

        replace_token = "[X]"
        mask_token = "[MASK]"
        mask_ids = 103
        prompt_templates = ['"{}"，它的意思是[MASK]。'.format(replace_token), '"{}"，这句话的意思是[MASK]。'.format(replace_token)]
        dropout = 0.15
        train_path = root_path + '/data/train.csv'
        valid_path = root_path + '/data/devdata.csv'
        log_dir = root_path + '/Log/'
        batch_size = 6
        lr = 5e-6
        num_epochs = 300
        require_improvement = 60

        save_path = root_path + '/UseModel/'

    config = Config()
    model = PromptBert(config).to(config.device)
    chechpoint = torch.load(config.save_path + '/train_loss_best.pth.tar')
    model.load_state_dict(chechpoint['state_dict'])
    # model = nn.DataParallel(model, device_ids=[0, 1])
    model.eval()
    model = model.cuda()

    input_names = ['input_ids', 'input_tem']
    output_names = ['output1']

    input1 = torch.tensor([101] + [random.randint(100, 10000) for i in range(120)] + [103, 102], dtype=torch.long,
                          device=config.device).view(1, 123)
    input2 = torch.tensor([101] + [random.randint(100, 10000) for i in range(120)] + [103, 102], dtype=torch.long,
                          device=config.device).view(1, 123)

    dummy_input = (input1, input2)
    with torch.no_grad():
        torch.onnx.export(model, dummy_input, "prompt-bert.onnx", opset_version=11, verbose=True,
                          input_names=input_names, output_names=output_names,
                          dynamic_axes={'input_ids': [0, 20], 'input_tem': [0, 20]})


def testonnx():
    path = '/data0/jianyu10/PTM/huggingface_model_cache/chinese-roberta-wwm-ext/vocab.txt'
    w = open(path, 'r', encoding='utf-8').readlines()
    word2idx = {w.strip(): i for i, w in enumerate(w)}
    sess_options = onnxruntime.SessionOptions()
    export_model_path = 'prompt-bert.onnx'
    session = onnxruntime.InferenceSession(export_model_path, sess_options, providers=['CPUExecutionProvider'])
    titles = ['英国首相府违规聚会调查报告认为政府“领导不力”---调查报告认为政府“领导不力”', '在欢声笑语中展现新时代新征程上精气神---2022年春节联欢晚会——在欢声笑语中展现新时代新征程上精气神',
              '因发生斗殴事件致2名犯人死亡 美国联邦监狱进入封锁状态---美国休斯敦发生枪击事件 致1名警员死亡', '英媒：泽连斯基下令三年扩军十万 敦促议员不要散布恐慌---泽连斯基签署法令：未来三年内扩军十万人',
              '冬季风暴渐平息 美部分地区降雪超50厘米---冬季风暴持续影响美国大部分地区', '英媒：特朗普大举筹款瞄准2024总统大选---特朗普又夸下海口：若2024再当选总统，将赦免国会大厦骚乱者',
              '多国政要和国际组织官员贺新春 祝新年如虎添翼---视频｜多位国际组织负责人及国家政要送上新春祝福', '“美国‘超额死亡’人数已近百万”---华尔街日报：疫情下，美国“超额死亡”人数已近百万',
              '日本外相：驻日美军入境新冠检测所用方法未被日本认可有效---外媒：驻日美军人员被曝离美前未进行新冠检测',
              '除夕，布林肯又发新闻公报拜年：愿虎年给所有人带来机遇、成功和健康---美国务卿布林肯拜年：愿虎年给所有人带来机遇、成功和健康']

    def getids(title):
        ori = title + '，它的意思是'
        tem = title + '，这句话的意思是'
        oriids = [101] + [word2idx.get(i, 100) for i in ori] + [103] + [word2idx.get('。', 100)] + [102]
        temids = [101] + [word2idx.get(i, 100) for i in tem] + [103] + [word2idx.get('。', 100)] + [102]
        oriids = oriids + [0] * (123 - len(oriids))
        temids = temids + [0] * (123 - len(temids))
        return oriids, temids

    for corpus in titles:
        title = corpus.split('---')
        orititle = title[0]
        simtitle = title[1]
        oriids, temids = getids(orititle)
        simids, simtemids = getids(simtitle)
        start = time.time()
        ort_inputs = {
            'input_ids': np.array([oriids]),
            'input_tem': np.array([temids]),
        }
        ort_outputs = session.run(['output1'], ort_inputs)
        outputs = torch.tensor(ort_outputs[0])
        cost = (time.time() - start) * 1000
        print(cost, outputs)
    # print("OnnxRuntime cpu Inference time = {} ms".format(format(sum(latency) * 1000 / len(latency), '.2f')))


def onnx2trt():
    pass


if __name__ == '__main__':
    testonnx()
